package org.strieber.fca.mid3;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.Stack;

import org.strieber.fca.Rule;

public class ID3 {

	// contains the possible values for each attribute 
	HashMap<String, List<String>> domains = new HashMap<String, List<String>>();
	ArrayList<String> usedAttributes = new ArrayList<String>();
	ArrayList<Rule> rules = new ArrayList<Rule>();
	DecisionNode root;
	
	public ID3() {
		
	}
	
	public ID3(HashMap<String, List<String>> domains, ArrayList<Rule> rules) {
		this.domains = domains;
		this.rules = rules;
		String highestGainAttribute = findHighestGain(this.rules);
		DecisionNode node = new DecisionNode();
		node.setRules(rules);
		node.setEntropy(calculateMainEntropy());
		node.setAttributeSplitOn(highestGainAttribute);
		root = node;
		expandNode(root);
	}
	
	protected void expandNode(DecisionNode node) {
		// grab the node we are working on
		DecisionNode currentNode = node;
		// get all the rules from that node
		List<Rule> nodeRules = currentNode.getRules();
		// if the rule list is empty or if all the intents are the same, skip it
		if ((nodeRules.size() != 0) && !allSameDetermined(nodeRules)) {
			// find the attribute the currentNode will be split on
			String attributeSplitOn = currentNode.getAttributeSplitOn();
			// get all the possible values for that attribute
			List<String> attributeValues = domains.get(attributeSplitOn);
			// add the attribute to the list of used attributes
			usedAttributes.add(attributeSplitOn);
			for (int i = 0; i < attributeValues.size(); i++) {
				// select one of the possible attribute values
				String currentAttributeValue = attributeValues.get(i);
				// find all the matching rules of that attribute/value pair
				List<Rule> matchingRules = findMatchingRules(attributeSplitOn, currentAttributeValue, nodeRules);
				
				if (matchingRules.size() > 0) {
					// create a new child node to attach to the current node
					DecisionNode childNode = new DecisionNode();
					// set the rules of the child as the matching rules
					childNode.setRules(matchingRules);
					// find the next highest gain attribute
					String highestGainAttribute = findHighestGain(matchingRules);
					// set the child node's highest gain attribute
					childNode.setAttributeSplitOn(highestGainAttribute);
					// set the child's value split on
					childNode.setValueSplitOn(currentAttributeValue);
					// set the child's entropy
					childNode.setEntropy(calculateEntropy(matchingRules));
					// set the child's parent
					childNode.setParent(currentNode);
					// add the child to the parent
					currentNode.addChild(childNode);
				}			
			}
			// expand all of the current node's new children
			for (int j = 0; j < currentNode.getChildren().size(); j++) {
				expandNode(currentNode.getChildren().get(j));
			}
		}
	}
	
	protected boolean allSameDetermined(List<Rule> rules) {
		// grab the first rule's intent
		String determinedRule = rules.get(0).getIntents().get(0);
		boolean allSame = true;
		// compare the first rule's intent to the rest
		for (int i = 0; i < rules.size(); i++) {
			if (!determinedRule.equals(rules.get(i).getIntents().get(0))) {
				allSame = false;
			}
		}
		return allSame;
	}
	
	public void setDomains(HashMap<String, List<String>> domains) {
		this.domains = domains;
	}
	
	public void setRules(ArrayList<Rule> rules) {
		this.rules = rules;
	}
	
	protected String findHighestGain(List<Rule> ruleSet) {
		String determinedAttribute = ruleSet.get(0).getIntents().get(0).split("=")[0];
		Set<String> attributes = domains.keySet();
		String highestAttribute = "";
		double highestGain = 0;
		
		Iterator<String> it = attributes.iterator();
		
		while (it.hasNext()) {
			String attribute = it.next();
			if (!attribute.equals(determinedAttribute)) {
				double gain = calculateGain(ruleSet, attribute);
				// if the attribute has been split on already, do not consider it
				if (gain > highestGain && !usedAttributes.contains(attribute)) {
					highestGain = gain;
					highestAttribute = attribute;
				}
			}
		}
		return highestAttribute;
	}
	
	protected double calculateMainEntropy() {
		return calculateEntropy(this.rules);
	}
	
	protected double calculateEntropy(List<Rule> ruleSet) {
		int numberOfRules = ruleSet.size();
		if (numberOfRules == 0) {
			return 0;
		}
		// assuming only one determinedAttribute for now
		String determinedAttribute = ruleSet.get(0).getIntents().get(0).split("=")[0];
		List<String> determinedValues = domains.get(determinedAttribute);
		double sum = 0;
		for (int i = 0; i < determinedValues.size(); i++) {
			int count = 0;
			String currentValue = determinedValues.get(i);
			for (int j = 0; j < numberOfRules; j++) {
				String ruleValue = ruleSet.get(j).getIntents().get(0).split("=")[1];
				// check and see how many times this determinedAttribute value is present in the list
				if (ruleValue.equals(currentValue)) {
					count++;
				}
			}
			double probability = 1.0 * count / numberOfRules;
			if (count > 0) {
				double logOfProbability = Math.log(probability) / Math.log(2);
				sum += -probability * logOfProbability;
			}
		}
		return sum;
	}
	
	protected double calculateGain(List<Rule> rules, String splitAttribute) {	
		int numberOfRules = rules.size();
		// should i do something if rules are empty?
		
		double entropyOfSet = calculateMainEntropy();
		double entropyOfAttribute = 0;
		
		List<String> splitValues = domains.get(splitAttribute);
		for (int i = 0; i < splitValues.size(); i++) {
			String attributeValue = splitValues.get(i);
			ArrayList<Rule> matchingRules = new ArrayList<Rule>();
			for (int j = 0; j < numberOfRules; j++) {
				Rule currentRule = rules.get(j);
				String extentToMatch = splitAttribute + "=" + attributeValue;
				if (currentRule.hasExtent(extentToMatch)) {
					matchingRules.add(currentRule);
				}
			}
			entropyOfAttribute += (1.0 * matchingRules.size() / rules.size()) * calculateEntropy(matchingRules);
		}
		return entropyOfSet - entropyOfAttribute;
	}
	
	protected List<Rule> findMatchingRules(String attributeName, String attributeValue, List<Rule> targetRules) {
		//System.out.println("fMR: aName: " + attributeName + " aValue: " + attributeValue + " tRules: " + targetRules);
		ArrayList<Rule> matchingRules = new ArrayList<Rule>();
		String attributeToMatch = attributeName + "=" + attributeValue;
		//System.out.println("fMR: attributeToMatch: " + attributeToMatch);
		for (int i = 0; i < targetRules.size(); i++) {
			Rule currentRule = targetRules.get(i);
			if (currentRule.hasExtent(attributeToMatch)) {
				//System.out.println("fMR: matched! " + attributeToMatch);
				matchingRules.add(currentRule);
			}
		}
		return matchingRules;
	}
	
	public String toString() {
		StringBuilder sb = new StringBuilder();
		Stack<DecisionNode> stack = new Stack<DecisionNode>();
		stack.push(root);
		while (!stack.isEmpty()) {
			DecisionNode currentNode = stack.pop();
			sb.append(currentNode + "\n");
			if (currentNode.getChildren() != null) {
				for (int i = 0; i < currentNode.getChildren().size(); i++) {
					stack.push(currentNode.getChildren().get(i));
				}
			}
		}
		return sb.toString();
	}
	
	public static void main(String[] args) {
		// setup domains
		HashMap<String, List<String>> domains = new HashMap<String, List<String>>();
		ArrayList<String> values = new ArrayList<String>();
		
		values.add("short");
		values.add("tall");
		domains.put("Height", values);
		
		values = new ArrayList<String>();
		values.add("blonde");
		values.add("red");
		values.add("dark");
		domains.put("Hair", values);
		
		values = new ArrayList<String>();
		values.add("blue");
		values.add("brown");
		domains.put("Eye", values);
		
		values = new ArrayList<String>();
		values.add("+");
		values.add("-");
		domains.put("Class", values);
		
		ArrayList<Rule> rules = new ArrayList<Rule>();
		ArrayList<String> extents = new ArrayList<String>();
		ArrayList<String> intents = new ArrayList<String>();
		
		Rule tempRule = new Rule();
		
		double support = 50;
		double confidence = 75;
		
		// object 1
		extents.add("Height=short");
		extents.add("Hair=blonde");
		extents.add("Eye=blue");
		intents.add("Class=+");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 2
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();
		extents.add("Height=tall");
		extents.add("Hair=blonde");
		extents.add("Eye=brown");
		intents.add("Class=-");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 3
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();
		extents.add("Height=tall");
		extents.add("Hair=red");
		extents.add("Eye=blue");
		intents.add("Class=+");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 4
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();
		extents.add("Height=short");
		extents.add("Hair=dark");
		extents.add("Eye=blue");
		intents.add("Class=-");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 5
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();
		extents.add("Height=tall");
		extents.add("Hair=dark");
		extents.add("Eye=blue");
		intents.add("Class=-");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 6
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();
		extents.add("Height=tall");
		extents.add("Hair=blonde");
		extents.add("Eye=blue");
		intents.add("Class=+");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 7
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();
		extents.add("Height=tall");
		extents.add("Hair=dark");
		extents.add("Eye=brown");
		intents.add("Class=-");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// object 8
		extents = new ArrayList<String>();
		intents = new ArrayList<String>();		
		extents.add("Height=short");
		extents.add("Hair=blonde");
		extents.add("Eye=brown");
		intents.add("Class=-");
		tempRule = new Rule(extents, intents, support, confidence);
		rules.add(tempRule);
		
		// create id3 object
		ID3 id3 = new ID3(domains, rules);
		System.out.println(id3);
	}
}